"""Async LanceDB database manager with connection pooling."""

import asyncio
import json
import os
import shutil
import signal
import sys
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Dict, List, Optional, Any, AsyncGenerator
import logging

import lancedb
import numpy as np
import pyarrow as pa
from lancedb import AsyncConnection, AsyncTable

logger = logging.getLogger(__name__)


class DBConnection:
    """Async database connection wrapper."""

    def __init__(self, uri: str):
        self.uri = uri
        self.connection: Optional[AsyncConnection] = None
        self.in_use = False
        self.last_used = datetime.now()
        self._lock = asyncio.Lock()

    async def connect(self) -> AsyncConnection:
        """Async connection initialization."""
        async with self._lock:
            if self.connection:
                return self.connection
            try:
                self.connection = await lancedb.connect_async(self.uri)
                self.last_used = datetime.now()
                logger.debug(f"Connected to LanceDB at {self.uri}")
                return self.connection
            except Exception as e:
                logger.error(f"Failed to connect to {self.uri}: {e}")
                raise ConnectionError(f"Failed to connect: {str(e)}")

    async def close(self):
        """Async close connection."""
        async with self._lock:
            if self.connection is None:
                return

            try:
                await self.connection.close()
                self.connection = None
                self.in_use = False
                logger.debug(f"Closed connection to {self.uri}")
            except Exception as e:
                logger.error(f"Failed to close connection: {e}")
            finally:
                self.connection = None

    @property
    def is_closed(self) -> bool:
        """Check if connection is closed."""
        return self.connection is None

    def __repr__(self):
        return f"DBConnection(uri={self.uri}, in_use={self.in_use}, closed={self.is_closed})"


class ConnectionPool:
    """Async database connection pool."""

    def __init__(self, uri: str, max_size: int = 10, min_size: int = 2):
        self.uri = uri
        self.max_size = max_size
        self.min_size = min_size
        self.pool: List[DBConnection] = []
        self._lock = asyncio.Lock()
        self._semaphore = asyncio.Semaphore(max_size)
        self._housekeeper_task: Optional[asyncio.Task] = None

    async def initialize(self):
        """Initialize connection pool."""
        async with self._lock:
            for _ in range(self.min_size):
                conn = DBConnection(self.uri)
                await conn.connect()
                self.pool.append(conn)
        logger.info(f"Initialized connection pool with {self.min_size} connections")

    async def start_housekeeper(self):
        """Start connection housekeeping task."""
        self._housekeeper_task = asyncio.create_task(self._housekeeper())
        logger.info("Started connection pool housekeeper")

    async def _housekeeper(self):
        """Connection housekeeping task."""
        while True:
            try:
                await asyncio.sleep(60)  # Check every minute
                async with self._lock:
                    now = datetime.now()
                    for conn in list(self.pool):
                        # Close idle connections older than 30 minutes
                        if not conn.in_use and (now - conn.last_used).seconds > 1800:
                            await conn.close()
                            self.pool.remove(conn)
                            logger.info(f"Closed idle connection (unused for 30+ minutes)")
            except Exception as e:
                logger.error(f"Housekeeper error: {e}")

    async def stop_housekeeper(self):
        """Stop connection housekeeping task."""
        if self._housekeeper_task:
            self._housekeeper_task.cancel()
            try:
                await self._housekeeper_task
            except asyncio.CancelledError:
                pass
            logger.info("Stopped connection pool housekeeper")

    @property
    def active_connections(self) -> int:
        """Current active connections."""
        return sum(1 for conn in self.pool if conn.in_use)

    @property
    def total_connections(self) -> int:
        """Total connections (active + idle)."""
        return len(self.pool)

    @property
    def idle_connections(self) -> int:
        """Idle connections."""
        return sum(1 for conn in self.pool if not conn.in_use)

    async def acquire(self) -> DBConnection:
        """Acquire connection."""
        await self._semaphore.acquire()

        async with self._lock:
            # Reuse idle connection
            for conn in self.pool:
                if not conn.in_use:
                    conn.in_use = True
                    conn.last_used = datetime.now()
                    return conn

            # Create new connection if under limit
            if len(self.pool) < self.max_size:
                conn = DBConnection(self.uri)
                await conn.connect()
                conn.in_use = True
                conn.last_used = datetime.now()
                self.pool.append(conn)
                return conn

            # Wait for available connection
            while True:
                for conn in self.pool:
                    if not conn.in_use:
                        conn.in_use = True
                        conn.last_used = datetime.now()
                        return conn
                await asyncio.sleep(0.1)

    async def release(self, connection: DBConnection):
        """Release connection."""
        connection.in_use = False
        self._semaphore.release()

    async def close_all(self):
        """Close all connections."""
        async with self._lock:
            connections = list(self.pool)
            self.pool.clear()

            for conn in connections:
                try:
                    await conn.close()
                except Exception as e:
                    logger.error(f"Error closing connection: {e}")

        await self.stop_housekeeper()

    def get_stats(self) -> Dict[str, Any]:
        """Get pool statistics."""
        return {
            "total_connections": self.total_connections,
            "active_connections": self.active_connections,
            "idle_connections": self.idle_connections,
            "max_size": self.max_size,
            "min_size": self.min_size
        }


class AsyncVectorDBManager:
    """Async LanceDB manager with connection pooling."""

    def __init__(self, uri: str = "./data/vectordb", config_path: str = "./data/vectordb_config.json",
                 max_connections: int = 10, min_connections: int = 2):
        self.uri = uri
        self.config_path = config_path
        self.pool = ConnectionPool(uri, max_connections, min_connections)
        self.tables_info: Dict[str, Dict] = {}
        self.backup_path = f"{uri}_backup"
        self.last_error: Optional[str] = None
        self._initialized = False
        self._cleanup_done = False

        # Set up signal handlers
        if sys.platform != "win32":
            signal.signal(signal.SIGINT, self._signal_handler)
            signal.signal(signal.SIGTERM, self._signal_handler)

    def _signal_handler(self, signum, frame):
        """Handle process signals."""
        logger.info(f"Received signal {signum}. Performing cleanup...")
        asyncio.create_task(self.cleanup())

    async def initialize(self):
        """Initialize the manager."""
        if self._initialized:
            return

        try:
            # Ensure data directory exists
            os.makedirs(os.path.dirname(self.uri), exist_ok=True)

            # Initialize connection pool
            await self.pool.initialize()
            await self.pool.start_housekeeper()

            # Load existing state
            self.load_state()

            self._initialized = True
            logger.info(f"Initialized AsyncVectorDBManager with URI: {self.uri}")

        except Exception as e:
            logger.error(f"Failed to initialize AsyncVectorDBManager: {e}")
            raise

    async def cleanup(self):
        """Cleanup resources."""
        if self._cleanup_done:
            return

        try:
            logger.info("Starting AsyncVectorDBManager cleanup...")
            self.save_state()
            await self.pool.close_all()
            self._cleanup_done = True
            logger.info("AsyncVectorDBManager cleanup completed")
        except Exception as e:
            logger.error(f"Error during cleanup: {e}")

    def save_state(self) -> bool:
        """Save database state to configuration file."""
        try:
            config = {
                "db_uri": self.uri,
                "tables_info": self.tables_info,
                "last_saved": datetime.now().isoformat(),
                "pool_stats": self.pool.get_stats()
            }

            # Create backup
            if os.path.exists(self.uri):
                if os.path.exists(self.backup_path):
                    shutil.rmtree(self.backup_path)
                shutil.copytree(self.uri, self.backup_path)
                logger.info(f"Created backup at {self.backup_path}")

            # Save configuration
            os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
            with open(self.config_path, 'w') as f:
                json.dump(config, f, indent=2)

            logger.info(f"Saved state to {self.config_path}")
            return True

        except Exception as e:
            self.last_error = str(e)
            logger.error(f"Failed to save state: {e}")
            return False

    def load_state(self) -> bool:
        """Load database state from configuration file."""
        try:
            if os.path.exists(self.config_path):
                with open(self.config_path, 'r') as f:
                    config = json.load(f)

                self.uri = config.get("db_uri", self.uri)
                self.tables_info = config.get("tables_info", {})

                logger.info(f"Loaded state from {self.config_path}")
                return True
            return False

        except Exception as e:
            self.last_error = str(e)
            logger.error(f"Failed to load state: {e}")
            return False

    @asynccontextmanager
    async def get_connection(self) -> AsyncGenerator[AsyncConnection, None]:
        """Get database connection context manager."""
        conn = await self.pool.acquire()
        try:
            yield conn.connection
        finally:
            await self.pool.release(conn)

    async def create_table(self, table_name: str, schema: Optional[pa.Schema] = None,
                          vector_dim: int = 1024, overwrite: bool = False) -> AsyncTable:
        """Create a new table."""
        if schema is None:
            schema = pa.schema([
                pa.field("vector", pa.list_(pa.float32(), vector_dim)),
                pa.field("text", pa.string()),
                pa.field("metadata", pa.string()),
                pa.field("created_at", pa.timestamp('us'))
            ])

        async with self.get_connection() as conn:
            try:
                if overwrite:
                    table = await conn.create_table(table_name, schema=schema, mode="overwrite")
                else:
                    table = await conn.create_table(table_name, schema=schema, exist_ok=True)

                self.tables_info[table_name] = {
                    "schema": schema.to_string(),
                    "created_at": datetime.now().isoformat(),
                    "vector_dim": vector_dim
                }

                logger.info(f"Created table '{table_name}' with vector dimension {vector_dim}")
                return table

            except Exception as e:
                self.last_error = str(e)
                logger.error(f"Failed to create table '{table_name}': {e}")
                raise

    async def search(self, table_name: str, query_vector: List[float], limit: int = 10,
                    threshold: float = 0.0) -> List[Dict[str, Any]]:
        """Perform vector search."""
        async with self.get_connection() as conn:
            try:
                table = await conn.open_table(table_name)

                # Perform vector search
                query = table.vector_search(query_vector).limit(limit)
                results = await query.to_list()

                # Apply threshold filtering if specified
                if threshold > 0:
                    results = [r for r in results if r.get('_distance', float('inf')) <= threshold]

                logger.debug(f"Search in table '{table_name}' returned {len(results)} results")
                return results

            except Exception as e:
                self.last_error = str(e)
                logger.error(f"Search failed in table '{table_name}': {e}")
                raise

    async def insert_data(self, table_name: str, data: List[Dict[str, Any]]) -> bool:
        """Insert data into table."""
        # Add timestamps if not present
        for item in data:
            if 'created_at' not in item:
                item['created_at'] = datetime.now()

        async with self.get_connection() as conn:
            try:
                table = await conn.open_table(table_name)
                await table.add(data)
                logger.info(f"Inserted {len(data)} records into table '{table_name}'")
                return True

            except Exception as e:
                self.last_error = str(e)
                logger.error(f"Failed to insert data into table '{table_name}': {e}")
                return False

    async def delete_data(self, table_name: str, conditions: Dict[str, Any]) -> bool:
        """Delete data from table based on conditions."""
        async with self.get_connection() as conn:
            try:
                table = await conn.open_table(table_name)

                # Build delete conditions
                if 'id' in conditions:
                    await table.delete(f"id = {conditions['id']}")
                elif 'metadata' in conditions:
                    # Handle metadata-based deletion
                    meta_conditions = []
                    for k, v in conditions['metadata'].items():
                        meta_conditions.append(f"metadata->'{k}' = '{v}'")

                    if meta_conditions:
                        await table.delete(" AND ".join(meta_conditions))

                logger.info(f"Deleted data from table '{table_name}' based on conditions: {conditions}")
                return True

            except Exception as e:
                self.last_error = str(e)
                logger.error(f"Failed to delete data from table '{table_name}': {e}")
                return False

    async def list_tables(self) -> List[str]:
        """List all tables."""
        async with self.get_connection() as conn:
            try:
                tables = await conn.table_names()
                return tables
            except Exception as e:
                self.last_error = str(e)
                logger.error(f"Failed to list tables: {e}")
                return []

    async def get_table_info(self, table_name: str) -> Optional[Dict[str, Any]]:
        """Get table information."""
        async with self.get_connection() as conn:
            try:
                table = await conn.open_table(table_name)
                schema = table.schema
                count = await table.count_rows()

                return {
                    "name": table_name,
                    "schema": str(schema),
                    "row_count": count,
                    "info": self.tables_info.get(table_name, {})
                }

            except Exception as e:
                self.last_error = str(e)
                logger.error(f"Failed to get table info for '{table_name}': {e}")
                return None

    async def get_database_stats(self) -> Dict[str, Any]:
        """Get database statistics."""
        try:
            tables = await self.list_tables()
            table_stats = {}

            for table_name in tables:
                info = await self.get_table_info(table_name)
                if info:
                    table_stats[table_name] = {
                        "row_count": info["row_count"],
                        "created_at": self.tables_info.get(table_name, {}).get("created_at")
                    }

            return {
                "uri": self.uri,
                "tables": table_stats,
                "total_tables": len(tables),
                "pool_stats": self.pool.get_stats(),
                "last_saved": self.tables_info.get("last_saved")
            }

        except Exception as e:
            logger.error(f"Failed to get database stats: {e}")
            return {"error": str(e)}

    def get_last_error(self) -> Optional[str]:
        """Get last error message."""
        return self.last_error


# Global database manager instance
db_manager: Optional[AsyncVectorDBManager] = None


async def get_db_manager() -> AsyncVectorDBManager:
    """Get global database manager instance."""
    global db_manager
    if db_manager is None:
        from .config import get_database_config
        db_config = get_database_config()
        db_manager = AsyncVectorDBManager(
            uri=db_config.uri,
            config_path=db_config.config_path,
            max_connections=db_config.max_connections,
            min_connections=db_config.min_connections
        )
        await db_manager.initialize()
    return db_manager


async def close_db_manager():
    """Close global database manager."""
    global db_manager
    if db_manager:
        await db_manager.cleanup()
        db_manager = None